from confluent_kafka import Producer, Consumer
from confluent_kafka.admin import AdminClient, NewTopic


class KafkaProducer:
    def __init__(self, **kwargs):
        self.linger_ms = 500
        self.batch_size = 1000
        self.acks = 1
        self.producer = None

        if kwargs:
            self.init(**kwargs)

    def init(self, **kwargs):
        self.producer = Producer({'linger.ms': self.linger_ms, 'batch.size': self.batch_size, 'acks': self.acks,
                                  **kwargs})

    def send(self, topic, value, *args, **kwargs):
        try:
            self.producer.produce(topic, value, *args, **kwargs)
        except BufferError:
            self.producer.flush()


class KafkaConsumer:
    def __init__(self, **kwargs):
        self.consumer = None
        if kwargs:
            self.init(**kwargs)

    def init(self, **kwargs):
        self.consumer = Consumer(kwargs)

    def poll(self, timeout=1):
        return self.consumer.poll(timeout)

    def consume(self, num_messages=1, timeout=-1):
        return self.consumer.consume(num_messages=num_messages, timeout=timeout)


class KafkaAdmin:
    def __init__(self, **kwargs):
        self.admin = None
        if kwargs:
            self.init(**kwargs)

    def init(self, **kwargs):
        self.admin = AdminClient(kwargs)

    def list_topics(self):
        return self.admin.list_topics()

    def create_topics(self, topics):
        if isinstance(topics, str):
            topics = [topics]
        return self.admin.create_topics([NewTopic(item) for item in topics])

    def delete_topics(self, topics):
        if isinstance(topics, str):
            topics = [topics]
        return self.admin.delete_topics(topics)


kafka_producer = KafkaProducer()
kafka_consumer = KafkaConsumer()
kafka_admin = KafkaAdmin()


if __name__ == '__main__':
    kafka_admin.init(**{'bootstrap.servers': '192.168.101.29:9092'})
    kafka_producer.init(**{'bootstrap.servers': '192.168.101.29:9092'})
    kafka_consumer.init(**{'bootstrap.servers': '192.168.101.29:9092', 'group.id': 'b',
                           'default.topic.config': {'auto.offset.reset': 'smallest'}})
    kafka_consumer.consumer.subscribe(['a10'])


    def test_admin():
        result = kafka_admin.admin.list_topics()
        print(result.topics)

    def test_sender():
        kafka_producer.send('a10', '11')

    def test_consumer():
        message = kafka_consumer.poll(1)
        messages = kafka_consumer.consume()


    test_admin()
    test_sender()
    test_consumer()
